package com.smp.dao.impl;

import com.smp.constant.EsIndexConstant;
import com.smp.dao.IUserRepository;
import com.smp.entry.UserBean;
import com.smp.vo.PageInfo;
import com.smp.vo.UserAggretion;
import com.smp.vo.UserParam;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.BucketOrder;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.sort.SortBuilders;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.SearchHits;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.data.elasticsearch.core.query.UpdateQuery;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: xiaoyu.fang
 * @date: 2020/8/25
 * @version: v1.0
 * @description:
 **/
@Component
public class UserRepositoryImpl implements IUserRepository<UserBean, String> {

    @Autowired
    private ElasticsearchRestTemplate elasticsearchRestTemplate;


    @Override
    public void update(UserBean userBean) {
        Map<String, Object> updateMap = new HashMap<>(2);
        updateMap.put("name", userBean.getName());
        updateMap.put("sex", userBean.getSex());
        elasticsearchRestTemplate.update(UpdateQuery.builder(userBean.getId())
                .withDocument(Document.from(updateMap)).build(), IndexCoordinates.of(EsIndexConstant.USER_BEAN));
    }

    @Override
    public PageInfo<UserBean> query(UserParam userParam) {
        NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder();
        if (!StringUtils.isEmpty(userParam.getName())) {
            nativeSearchQueryBuilder.withQuery(QueryBuilders.termQuery("name", userParam.getName()));
        }
        NativeSearchQuery nativeSearchQuery = nativeSearchQueryBuilder
                .withPageable(PageRequest.of(userParam.getPageNum() - 1, userParam.getPageSize()))
                .withSort(SortBuilders.fieldSort("createAt").order(SortOrder.DESC))
                .build();

        SearchHits<UserBean> searchHits = elasticsearchRestTemplate.search(nativeSearchQuery, UserBean.class);
        List<UserBean> userBeans = new ArrayList<>(searchHits.getSearchHits().size());
        searchHits.get().forEach(userBeanSearchHit -> userBeans.add(userBeanSearchHit.getContent()));
        return new PageInfo<>(userParam.getPageNum(), userParam.getPageSize(), (int) searchHits.getTotalHits(), userBeans);
    }

    @Override
    public List<UserBean> queryUserBeanByBirthDay(String beginAt) {
        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder().withQuery(QueryBuilders.rangeQuery("birthDay").gt(beginAt)).build();
        SearchHits<UserBean> searchHits = elasticsearchRestTemplate.search(searchQuery, UserBean.class);
        List<UserBean> userBeans = new ArrayList<>(searchHits.getSearchHits().size());
        searchHits.forEach(userBeanSearchHit -> userBeans.add(userBeanSearchHit.getContent()));
        return userBeans;
    }

    @Override
    public UserAggretion queryByAggregations() {
        UserAggretion userAggretion = new UserAggretion();
        QueryBuilder matchAllQuery = QueryBuilders.matchAllQuery();
        NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder();
        nativeSearchQueryBuilder.withQuery(matchAllQuery).withPageable(PageRequest.of(0, 1));
        // 按sex分组
        TermsAggregationBuilder termsAggregation = AggregationBuilders.terms("性别").field("sex").order(BucketOrder.key(true));
        nativeSearchQueryBuilder.addAggregation(termsAggregation);

        SearchHits<UserBean> searchHits = elasticsearchRestTemplate.search(nativeSearchQueryBuilder.build(), UserBean.class);
        Terms terms = (Terms) searchHits.getAggregations().asMap().get("性别");
        List<Terms.Bucket> buckets = (List<Terms.Bucket>) terms.getBuckets();
        Map<String, Object> stats = new HashMap<>(2);
        buckets.stream().forEach(bucket -> stats.put("1".equals(bucket.getKeyAsString()) ? "男" : "女", bucket.getDocCount()));
        userAggretion.setName(terms.getName());
        userAggretion.setStats(stats);
        return userAggretion;
    }
}
